package com.xh.common.thread.ext.limit;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * 执行限制任务的线程池执行器，实现了同一组任务能够指定同时执行的数量限制。
 */
public class LimitThreadPoolExecutor<T> extends AdapterExecutorService implements LimitExecutorService<T> {

    private ExecutorFilter<T> mFilter;
    private Map<String, LimitWorker<T>> mWorkers;
    private ReadWriteLock mLock;

    public LimitThreadPoolExecutor(ExecutorService service) {
        this(service, null);
    }

    public LimitThreadPoolExecutor(ExecutorService service, ExecutorFilter<T> filter) {
        super(service);
        mWorkers = new HashMap<>();
        mLock = new ReentrantReadWriteLock();

        setExecutorFilter(filter);
    }

    @Override
    public void setExecutorFilter(ExecutorFilter<T> filter) {
        mFilter = filter;
    }

    @Override
    public void execute(Runnable cmd) {
        assertRunnableAvailable(cmd);
        executeInternal(cmd, null);
    }

    @Override
    public void execute(Runnable cmd, T t) {
        assertRunnableAvailable(cmd);
        executeInternal(cmd, t);
    }

    @Override
    public void execute(Runnable cmd, String groupID, int limitSize) {
        execute(cmd, groupID, limitSize, null);
    }

    @Override
    public void execute(Runnable cmd, String groupID, int limitSize, T o) {
        assertRunnableAvailable(cmd);

        if (isShutdown()) {
            return;
        }

        if (groupID == null || groupID.isEmpty()) {
            executeInternal(cmd, o);
            return;
        }

        LimitWorker<T> worker = checkAndGetWorker(groupID, limitSize);

        if (limitSize <= 0) {
            executeInternal(cmd, o);
        } else {
            worker.execute(cmd, o);
        }
    }

    private LimitWorker<T> checkAndGetWorker(String groupID, int limitSize) {
        LimitWorker<T> worker;

        // 先用读锁来获取
        Lock readLock = mLock.readLock();
        readLock.lock();
        try {
            worker = mWorkers.get(groupID);
        } finally {
            readLock.unlock();
        }

        if (worker != null) {
            assertWorkerAvailable(worker, limitSize);
            return worker;
        }

        // 如果没有获取到，再用写锁来获取
        Lock writeLock = mLock.writeLock();
        writeLock.lock();
        try {
            worker = mWorkers.get(groupID);
            if (worker != null) {
                assertWorkerAvailable(worker, limitSize);
            } else {
                worker = new LimitWorker<>(this, limitSize);
                mWorkers.put(groupID, worker);
            }
        } finally {
            writeLock.unlock();
        }
        return worker;
    }

    private void assertRunnableAvailable(Runnable runnable) {
        if (runnable == null) {
            throw new NullPointerException();
        }
    }

    private void assertWorkerAvailable(LimitWorker<T> worker, int limitSize) {
        if (worker != null && limitSize != worker.getLimitSize()) {
            throw new IllegalArgumentException();
        }
    }

    void executeInternal(Runnable command, T t) {
        if (isShutdown()) {
            return;
        }
        if (mFilter != null) {
            mFilter.execute(mService, command, t);
        } else {
            mService.execute(command);
        }
    }

    @Override
    public void shutdown() {
        mService.shutdown();
        Lock lock = mLock.writeLock();
        lock.lock();
        try {
            mWorkers.clear();
        } finally {
            lock.unlock();
        }
    }

    @Override
    public List<Runnable> shutdownNow() {
        List<Runnable> list = mService.shutdownNow();

        Lock lock = mLock.writeLock();
        lock.lock();
        try {
            for (LimitWorker<T> worker : mWorkers.values()) {
                list.addAll(worker.shutdownNow());
            }
            mWorkers.clear();
        } finally {
            lock.unlock();
        }

        return list;
    }
}
